/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com> 
 *
 */
#ifndef SRC_PROTOCOLS_UDP_UDPPROTOCOL_H_
#define SRC_PROTOCOLS_UDP_UDPPROTOCOL_H_

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include "Multiplexer.h"
#include "Protocol.h"
#include <netinet/udp.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include "flow/FlowManager.h"
#include "flow/FlowCache.h"
#include "FlowForwarder.h"
#include "DatabaseAdaptor.h"

namespace aiengine {

class UDPProtocol: public Protocol {
public:
	explicit UDPProtocol(const std::string& name):
		Protocol(name) {}
	explicit UDPProtocol():UDPProtocol(UDPProtocol::default_name) {}
    	virtual ~UDPProtocol();

	static constexpr const char *const default_name = "UDP";
	static const uint16_t header_size = 8;

	uint16_t getId() const override { return IPPROTO_UDP; }
	uint16_t getHeaderSize() const override { return header_size; }

	// Condition for say that a packet is UDP 
	bool check(const Packet &packet) override;
	void processFlow(Flow *flow) override {} // This protocol generates flows but not for destination.
	bool processPacket(Packet &packet) override;

	void statistics(std::basic_ostream<char> &out, int level, int32_t limit) const override;
	void statistics(Json &out, int level) const override;

        void releaseCache() override {} // No need to free cache

        void setHeader(const uint8_t *raw_packet) override {
        
                header_ = reinterpret_cast <const udphdr*> (raw_packet);
        }

#if defined(IS_FREEBSD) || defined(IS_OPENBSD) || defined(IS_DARWIN)
	uint16_t getSourcePort() const { return ntohs(header_->uh_sport); }
    	uint16_t getDestinationPort() const { return ntohs(header_->uh_dport); }
    	uint16_t getLength() const { return ntohs(header_->uh_ulen); }
    	unsigned int getPayloadLength() const { return ntohs(header_->uh_ulen) - sizeof(struct udphdr); }
#else
	uint16_t getSourcePort() const { return ntohs(header_->source); }
    	uint16_t getDestinationPort() const { return ntohs(header_->dest); }
    	uint16_t getLength() const { return ntohs(header_->len); }
    	unsigned int getPayloadLength() const { return ntohs(header_->len) - sizeof(udphdr); }
#endif
    	unsigned int getHeaderLength() const { return sizeof(struct udphdr); }
	uint8_t* getPayload() const { return (uint8_t*)header_ +getHeaderLength(); }

	void setFlowManager(FlowManagerPtr flow_mng) { flow_table_ = flow_mng;}
	FlowManagerPtr getFlowManager() const { return flow_table_; }
	void setFlowCache(FlowCachePtr flow_cache) { flow_cache_ = flow_cache;}
	FlowCachePtr getFlowCache() const { return flow_cache_;}

#if defined(HAVE_REJECT_FLOW)
	void addRejectFunction(std::function <void (Flow*)> reject) { reject_func_ = reject; }
#endif

	void setRegexManager(const SharedPointer<RegexManager> &rm) { rm_ = rm;}

	Flow *getCurrentFlow() const { return current_flow_;} // used just for testing pourposes

        void increaseAllocatedMemory(int value) override;
        void decreaseAllocatedMemory(int value) override;

	uint64_t getCurrentUseMemory() const override; 
	uint64_t getAllocatedMemory() const override;
	uint64_t getTotalAllocatedMemory() const override;

        void setDynamicAllocatedMemory(bool value) override; 
        bool isDynamicAllocatedMemory() const override;

	uint32_t getTotalCacheMisses() const override;
	uint32_t getTotalEvents() const override { return total_events_; }

	CounterMap getCounters() const override; 
	void resetCounters() override;

	void setAnomalyManager(SharedPointer<AnomalyManager> amng) override { anomaly_ = amng; }

private:
	SharedPointer<Flow> getFlow(const Packet &packet); 

	FlowManagerPtr flow_table_ = nullptr;
	FlowCachePtr flow_cache_ = nullptr;
	SharedPointer<RegexManager> rm_ = nullptr;
	Flow *current_flow_ = nullptr;
	const udphdr *header_ = nullptr;
	uint32_t total_events_ = 0;
	time_t last_timeout_ = 0;
	time_t packet_time_ = 0;

#if defined(HAVE_REJECT_FLOW)
	std::function <void (Flow*)> reject_func_ = [] (Flow*) {};
#endif
	SharedPointer<AnomalyManager> anomaly_ = nullptr;
};

typedef std::shared_ptr<UDPProtocol> UDPProtocolPtr;

} // namespace aiengine

#endif  // SRC_PROTOCOLS_UDP_UDPPROTOCOL_H_
